package zk.util;

import lombok.extern.slf4j.Slf4j;
import org.I0Itec.zkclient.IZkDataListener;
import org.I0Itec.zkclient.ZkClient;

import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;

/**
 * 使用zookeeper临时有序节点+注册监听实现分布式锁
 * <p>
 * 使用zookeeper实现分布式锁的算法流程，假设锁空间的根节点为/lock：
 * <p>
 * 1、客户端连接zookeeper，并在/lock下创建临时的且有序的子节点，第一个客户端对应的子节点为/lock/lock-0000000000，第二个为/lock/lock-0000000001，以此类推。
 * 2、客户端获取/lock下的子节点列表，判断自己创建的子节点是否为当前子节点列表中序号最小的子节点，如果是则认为获得锁，否则监听/lock的子节点变更消息，获得子节点变更通知后重复此步骤直至获得锁
 * 3、执行业务代码；
 * 4、完成业务流程后，删除对应的子节点释放锁。
 *
 * @author duchao
 */
@Slf4j
public class ZkLock {

    public final static String ROOT_PATH = "/lock";
    public ZkClient zkClient;
    private String lockName;
    private String eNodeName;
    private String eNodeNo;

    public ZkLock(String lockName) {
        zkClient = new ZkClient("182.92.202.68:8022", 5000);
        this.lockName = ROOT_PATH.concat("/").concat(lockName);
        if (!zkClient.exists(ROOT_PATH)) {
            zkClient.createPersistent(ROOT_PATH);
        }
        if (!zkClient.exists(this.lockName)) {
            zkClient.createPersistent(this.lockName);
        }
    }

    /**
     * 加锁
     */
    public void lock() {
        for (; ; ) {
            //创建临时节点
            if (eNodeName == null) {
                eNodeName = zkClient.createEphemeralSequential(lockName.concat("/"), System.currentTimeMillis());
                eNodeNo = eNodeName.substring(lockName.length() + 1);
            }

            //获取全部节点
            List<String> children = zkClient.getChildren(lockName);

            //检查是否持有锁，如果是序号最小则视为持有
            Collections.sort(children);
            if (eNodeNo.equals(children.get(0))) {
                log.info("{}节点获取分布式锁", eNodeNo);
                break;
            }

            //CountDownLatch latch = new CountDownLatch(1);
            //判断自己排第几个，如果自己没有获得锁，则要监听前一个节点
            int selfIndex = Collections.binarySearch(children, eNodeNo);
            if (selfIndex < 0) {
                throw new RuntimeException("节点" + eNodeNo + "没找到");
            }
            String prePathNo = children.get(selfIndex - 1);
            zkClient.subscribeDataChanges(lockName.concat("/").concat(prePathNo), new IZkDataListener() {
                @Override
                public void handleDataChange(String dataPath, Object data) throws Exception {
                }

                @Override
                public void handleDataDeleted(String dataPath) throws Exception {
                    log.info("{}节点被唤醒", eNodeNo);
                    //latch.countDown();
                }
            });

            try {
                log.info("{}节点等待锁...", eNodeNo);
                //latch.await();
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }
    }

    /**
     * 释放锁
     */
    public void unlock() {
        zkClient.delete(eNodeName);
        zkClient.close();
        log.info("{}节点释放了锁", eNodeName);
    }

    public static void main(String[] args) {
        String LOCK_PREFIX = "lock:";
//        ZkClient zkClient = new ZkClient("182.92.202.68:8022", 5000);
//        zkClient.deleteRecursive(ROOT_PATH);

        int threadCount = 5;
        CountDownLatch latch = new CountDownLatch(threadCount);
        for (int i = 0; i < threadCount; i++) {
            new Thread(() -> {
                try {
                    latch.await();
                    ZkLock zkLock = new ZkLock(LOCK_PREFIX + 1);
                    zkLock.lock();
                    compute();
                    zkLock.unlock();
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }).start();
            latch.countDown();
        }
    }

    private static int count = 0;

    /**
     * 模拟业务操作
     */
    public static void compute() {
        count++;
        log.info("计数：{}\n", count);
    }

}
